{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Molecular Dynamics\n",
    "\n",
    "Run molecular dynamics (MD) simulations using [GPUMD](https://github.com/brucefan1983/GPUMD)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating bulk atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo, Morph\n",
    "atoms = SymbolInfo('MoTaVW', 'bcc', 3).create_bulk_atoms((20, 20, 20))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running molecular dynamics simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_in = ['potential nep.txt', \n",
    "          'velocity 300', \n",
    "          'time_step 1', \n",
    "          'ensemble npt_scr 300 300 200 0 500 2000',\n",
    "          'dump_thermo 1000', \n",
    "          'dump_restart 30000', \n",
    "          'dump_exyz 10000',\n",
    "          'run 30000']\n",
    "Morph(atoms).gpumd('relax', run_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Deforming the simulation box."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_in = ['potential nep.txt', \n",
    "          'velocity 300', \n",
    "          'time_step 1',\n",
    "          'ensemble npt_scr 300 300 100 0 0 0 100 100 100 1000',\n",
    "          'run 30000', \n",
    "          'ensemble npt_scr 300 300 100 0 0 0 100 100 100 1000',\n",
    "          'deform 0.00001 0 0 1', \n",
    "          'dump_thermo 1000', \n",
    "          'dump_exyz 1000', \n",
    "          'dump_restart 10000',\n",
    "          'run 1000000']\n",
    "Morph(atoms).gpumd('deform', run_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulating the process of a crystallization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_in = ['potential nep.txt', \n",
    "          'velocity 2000', \n",
    "          'time_step 1', \n",
    "          'ensemble npt_scr 2000 2000 200 0 500 2000', \n",
    "          'dump_thermo 1000', \n",
    "          'dump_exyz 10000', \n",
    "          'dump_restart 10000', \n",
    "          'run 100000',\n",
    "          'ensemble npt_scr 2000 5000 200 0 500 2000',\n",
    "          'dump_thermo 1000', \n",
    "          'dump_exyz 100000', \n",
    "          'dump_restart 10000', \n",
    "          'run 10000000',\n",
    "          'ensemble npt_scr 4500 4500 200 0 500 2000',\n",
    "          'dump_thermo 1000', \n",
    "          'dump_exyz 10000', \n",
    "          'dump_restart 10000', \n",
    "          'run 100000',\n",
    "          'ensemble npt_scr 4500 1500 200 0 500 2000',\n",
    "          'dump_thermo 1000', \n",
    "          'dump_exyz 100000', \n",
    "          'dump_restart 10000', \n",
    "          'run 10000000']\n",
    "Morph(atoms).gpumd('crystallization', run_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculating the Melting point using two-phase coexistence method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_restart\n",
    "\n",
    "group = []\n",
    "for atom in atoms:\n",
    "    if atom.position[2] < atoms.cell[2, 2] / 2:\n",
    "        group.append(0)\n",
    "    else:\n",
    "        group.append(1)\n",
    "atoms.info['group'] = group\n",
    "\n",
    "run_in_1 = ['potential nep.txt', \n",
    "            'velocity 3000', \n",
    "            'time_step 1', \n",
    "            'ensemble npt_ber 3000 3000 200 0 500 2000', \n",
    "            'dump_exyz 10000', \n",
    "            'dump_thermo 1000',\n",
    "            'run 30000',\n",
    "            'ensemble heat_lan 3500 200 500 0 1',\n",
    "            'dump_exyz 10000',\n",
    "            'dump_thermo 1000',\n",
    "            'dump_restart 10000',\n",
    "            'run 1000000']\n",
    "\n",
    "Morph(atoms).gpumd('melting_point/relax', run_in_1)\n",
    "\n",
    "for Tm in range(3400, 3701, 100):\n",
    "    atoms = read_restart('melting_point/relax/restart.xyz')\n",
    "    run_in = ['potential nep.txt', \n",
    "             f'velocity {Tm}', \n",
    "              'time_step 1', \n",
    "             f'ensemble npt_ber {Tm} {Tm} 200 0 500 2000', \n",
    "              'dump_exyz 10000', \n",
    "              'dump_thermo 1000',\n",
    "              'run 30000']\n",
    "    Morph(atoms).gpumd(f'melting_point/{Tm}', run_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulating the radiation damage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_restart\n",
    "import numpy as np\n",
    "\n",
    "group = []\n",
    "thickness = 3 * 3\n",
    "for atom in atoms:\n",
    "    if atom.position[0] < thickness or atom.position[1] < thickness or atom.position[2] < thickness:\n",
    "        group.append(0)\n",
    "    elif atom.position[0] >= atoms.cell[0, 0] - thickness or atom.position[1] >= atoms.cell[1, 1] - thickness or atom.position[2] >= atoms.cell[2, 2] - thickness:\n",
    "        group.append(1)\n",
    "    else:\n",
    "        group.append(2)\n",
    "atoms.info['group'] = group\n",
    "\n",
    "run_in_1 = ['potential nep.txt',\n",
    "            'velocity 300', \n",
    "            'time_step 1', \n",
    "            'ensemble npt_scr 300 300 200 0 500 2000', \n",
    "            'dump_thermo 1000', \n",
    "            'dump_restart 30000', \n",
    "            'run 30000']\n",
    "\n",
    "run_in_2 = ['potential nep.txt', \n",
    "            'velocity 300', \n",
    "            'time_step 0', \n",
    "            'ensemble nve',\n",
    "            'dump_exyz 1', \n",
    "            'run 1',\n",
    "            'time_step 1 0.015', \n",
    "            'ensemble heat_nhc 300 200 0 0 1',\n",
    "            'compute 0 200 10 temperature', \n",
    "            'dump_restart 10000', \n",
    "            'dump_exyz 2000 1 1',\n",
    "            'run 70000']\n",
    "\n",
    "pka_energy = 1000 #eV\n",
    "direction = np.array([1, 3, 5]) \n",
    "\n",
    "Morph(atoms).gpumd('radiation/relax', run_in_1)\n",
    "atoms = read_restart('radiation/relax/restart.xyz')\n",
    "Morph(atoms).set_pka(pka_energy, direction)\n",
    "Morph(atoms).gpumd('radiation/cascade', run_in_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulating the overlapping cascades."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_restart\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "run_in_1 = ['potential nep.txt',\n",
    "            'velocity 300', \n",
    "            'time_step 1', \n",
    "            'ensemble npt_scr 300 300 200 0 500 2000', \n",
    "            'dump_thermo 1000', \n",
    "            'dump_restart 10000', \n",
    "            'run 30000']\n",
    "\n",
    "run_in_2 = ['potential nep.txt',\n",
    "            'velocity 300', \n",
    "            'time_step 1', \n",
    "            'ensemble npt_scr 300 300 200 0 500 2000', \n",
    "            'dump_thermo 1000', \n",
    "            'dump_restart 10000', \n",
    "            'run 10000']\n",
    "\n",
    "run_in_3 = ['potential nep.txt', \n",
    "            'velocity 300', \n",
    "            'time_step 0', \n",
    "            'ensemble nve',\n",
    "            'dump_exyz 1', \n",
    "            'run 1',\n",
    "            'time_step 1 0.015', \n",
    "            'ensemble heat_nhc 300 200 0 0 1',\n",
    "            'electron_stop ../../electron_stopping_fit.txt',\n",
    "            'compute 0 200 10 temperature', \n",
    "            'dump_restart 10000', \n",
    "            'dump_exyz 2000 1 1',\n",
    "            'run 40000']\n",
    "\n",
    "pka_energy = 10 #eV\n",
    "cascade_times = 2000\n",
    "directions = [np.array([np.sin(np.random.uniform(0, np.pi)) * np.cos(np.random.uniform(0, 2 * np.pi)),\n",
    "                        np.sin(np.random.uniform(0, np.pi)) * np.sin(np.random.uniform(0, 2 * np.pi)),\n",
    "                        np.cos(np.random.uniform(0, np.pi))]) for _ in range(cascade_times)]\n",
    "\n",
    "indexs = [random.randint(0, len(atoms) - 1) for _ in range(cascade_times)]\n",
    "\n",
    "## First time\n",
    "direction = directions[0]\n",
    "index = indexs[0]\n",
    "\n",
    "center = atoms.cell.diagonal() / 2\n",
    "diff = center - atoms[index].position\n",
    "for atom in atoms:\n",
    "    atom.position += diff\n",
    "\n",
    "for atom in atoms:\n",
    "    atom.position %= atoms.cell.diagonal()\n",
    "\n",
    "group = []\n",
    "thickness = 3.185 * 3\n",
    "for atom in atoms:\n",
    "    if atom.position[0] < thickness or atom.position[1] < thickness or atom.position[2] < thickness:\n",
    "        group.append(0)\n",
    "    elif atom.position[0] >= atoms.cell[0, 0] - thickness or atom.position[1] >= atoms.cell[1, 1] - thickness or atom.position[2] >= atoms.cell[2, 2] - thickness:\n",
    "        group.append(1)\n",
    "    else:\n",
    "        group.append(2)\n",
    "atoms.info['group'] = group\n",
    "\n",
    "Morph(atoms).gpumd('radiation0/relax', run_in_1)\n",
    "atoms = read_restart('radiation0/relax/restart.xyz')\n",
    "Morph(atoms).set_pka(pka_energy, direction, index)\n",
    "Morph(atoms).gpumd('radiation0/cascade', run_in_3)\n",
    "\n",
    "## Loops\n",
    "for i in range(1, cascade_times):\n",
    "    direction = directions[i]\n",
    "    index = indexs[i]\n",
    "    atoms = read_restart(f'radiation{i-1}/cascade/restart.xyz')\n",
    "\n",
    "    center = atoms.cell.diagonal() / 2\n",
    "    diff = center - atoms[index].position\n",
    "    for atom in atoms:\n",
    "        atom.position += diff\n",
    "\n",
    "    for atom in atoms:\n",
    "        atom.position %= atoms.cell.diagonal()\n",
    "\n",
    "    group = []\n",
    "    thickness = 3.185 * 3\n",
    "    for atom in atoms:\n",
    "        if atom.position[0] < thickness or atom.position[1] < thickness or atom.position[2] < thickness:\n",
    "            group.append(0)\n",
    "        elif atom.position[0] >= atoms.cell[0, 0] - thickness or atom.position[1] >= atoms.cell[1, 1] - thickness or atom.position[2] >= atoms.cell[2, 2] - thickness:\n",
    "            group.append(1)\n",
    "        else:\n",
    "            group.append(2)\n",
    "    atoms.info['group'] = group\n",
    "\n",
    "    Morph(atoms).gpumd(f'radiation{i}/relax', run_in_2)\n",
    "    atoms = read_restart(f'radiation{i}/relax/restart.xyz')\n",
    "    Morph(atoms).set_pka(pka_energy, direction, index)\n",
    "    Morph(atoms).gpumd(f'radiation{i}/cascade', run_in_3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the threshold displacement energy surface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo, Morph\n",
    "from ase.md.velocitydistribution import MaxwellBoltzmannDistribution\n",
    "from itertools import combinations_with_replacement\n",
    "from math import gcd\n",
    "import numpy as np\n",
    "\n",
    "atoms = SymbolInfo('W', 'bcc', 3.185).create_bulk_atoms(supercell=(12,12,16))\n",
    "group = []\n",
    "thickness = 3.185 * 1\n",
    "for atom in atoms:\n",
    "    if atom.position[0] < thickness or atom.position[1] < thickness or atom.position[2] < thickness:\n",
    "        group.append(0)\n",
    "    elif atom.position[0] >= atoms.cell[0, 0] - thickness or atom.position[1] >= atoms.cell[1, 1] - thickness or atom.position[2] >= atoms.cell[2, 2] - thickness:\n",
    "        group.append(1)\n",
    "    else:\n",
    "        group.append(2)\n",
    "atoms.info['group'] = group\n",
    "\n",
    "run_in =['potential nep.txt', \n",
    "         'velocity 300', \n",
    "         'time_step 0', \n",
    "         'ensemble nve',\n",
    "         'dump_exyz 1', \n",
    "         'run 1',\n",
    "         'time_step 1 0.015', \n",
    "         'ensemble heat_nhc 36 200 0 0 1',\n",
    "         'dump_exyz 100 1 1',\n",
    "         'run 5000']\n",
    "\n",
    "MaxwellBoltzmannDistribution(atoms, temperature_K=36)\n",
    "atoms.info['velocities'] = atoms.get_velocities()\n",
    "\n",
    "hmax = 5\n",
    "directions = set()\n",
    "for hkl in combinations_with_replacement(range(hmax + 1), 3):\n",
    "    hkl = tuple(np.array(np.array(hkl) / gcd(*hkl), dtype=int))\n",
    "    if sum(hkl) <= 0:\n",
    "        continue\n",
    "    directions.add(np.array(hkl))\n",
    "\n",
    "for direction in directions:\n",
    "    for pka_energy in range(40, 200, 20):\n",
    "        Morph(atoms).set_pka(pka_energy, direction)\n",
    "        direction_path = ''.join([str(i) for i in direction])\n",
    "        Morph(atoms).gpumd(f'TDE/{direction_path}/{pka_energy}', run_in_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the formation eneregy of interstitial clusters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo, Morph\n",
    "from wizard.io import read_xyz\n",
    "import numpy as np\n",
    "       \n",
    "burger = (1, 0, 0)\n",
    "Rcut = 10\n",
    "thickness = 2\n",
    "atoms = SymbolInfo('W', 'bcc', 3.185).create_bulk_atoms(supercell=(30, 30, 30))\n",
    "atoms_energy = -12.597\n",
    "center = atoms.get_center_of_mass()\n",
    "for atom in atoms:\n",
    "    vector = atom.position - center\n",
    "    proj = abs(vector @ burger) / np.linalg.norm(burger)\n",
    "    R = np.sqrt(max(np.dot(vector, vector) - np.dot(proj, proj), 0))\n",
    "    if  R < Rcut and proj < thickness:\n",
    "        Morph(atoms).create_self_interstitial_atom(burger, index = atom.index)\n",
    "Morph(atoms).gpumd('sia_cluster',['potential nep.txt', 'ensemble nve', 'time_step 0',\n",
    "                                  'minimize fire 1.0e-4 1000','dump_exyz 1','run 1'])\n",
    "frames = read_xyz('sia_cluster/dump.xyz')\n",
    "atoms = frames[-1]\n",
    "formation_energy = atoms.info['energy'] -  atoms_energy * len(atoms)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simulating by ASE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo\n",
    "from wizard.molecular_dynamics import MolecularDynamics\n",
    "from calorine.calculators import CPUNEP\n",
    "\n",
    "calc = CPUNEP('../Repository/Wnep2/nep.txt')\n",
    "atoms = SymbolInfo('W', 'bcc', 3).create_bulk_atoms()\n",
    "MolecularDynamics(atoms, calc).NPT(steps=1000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "undefined.undefined.undefined"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
